import cmd
import sys
import json
import asyncio
import ollama  # 依赖于 'ollama' 库

from common import *

# ----------------------------------------------------------------------------
# 1. 一个简单自然语言命令处理类
# ----------------------------------------------------------------------------
class NaturalLanguageProcessor:
    """
    一个调用本地Ollama模型将自然语言转换为结构化命令的处理器。
    """

    def __init__(self, model_name: str = "gemma3:4b"):
        """
        初始化处理器。

        Args:
            model_name (str): 要在Ollama中使用的模型名称。
                              例如 "gemma:latest" 或 "gemma3:4b"。
        """
        self.model_name = model_name
        # 预先生成工具描述，以便在提示中使用
        self.tools_prompt = self._get_tools_prompt()

    def _get_tools_prompt(self) -> str:
        """
        生成描述可用命令的提示字符串，用于引导LLM。
        这是非常关键的提示工程部分。
        """
        return """
        可用的命令如下:
        - 'set <key> <value>': 设置一个键值对。值可以包含空格。
        - 'get <key>': 获取一个键对应的值。
        - 'delete <key>': 删除一个键值对。
        - 'list': 列出所有存储的键值对，此命令没有参数。
        - 'unknown': 如果用户意图不明确或与任何可用命令都不匹配。
        """

    async def process(self, user_input: str) -> dict:
        """
        异步处理用户输入，并从Ollama获取结构化命令。

        Args:
            user_input (str): 用户的自然语言输入。

        Returns:
            dict: 一个包含 'command' 和 'arguments' 的字典。
                  例如: {'command': 'set', 'arguments': ['user:age', '30']}
        """
        if not user_input:
            return {'command': 'error', 'arguments': ['输入不能为空。']}

        system_prompt = f"""
        你是一个智能助手，你的唯一任务是将用户的自然语言请求转换为一个精确的JSON对象。
        这个JSON对象必须包含两个键：'command' 和 'arguments'。
        'command' 的值必须是以下之一：'set', 'get', 'delete', 'list', 'unknown'。
        'arguments' 的值必须是一个字符串列表。

        {self.tools_prompt}

        根据用户的请求，严格按照以下JSON格式进行响应，不要添加任何额外的解释或文本。

        例如:
        - 用户输入: "帮我把 'user:id' 设置成 '12345'"
          你的输出: {{"command": "set", "arguments": ["user:id", "12345"]}}
        - 用户输入: "查一下 'user:id' 的值"
          你的输出: {{"command": "get", "arguments": ["user:id"]}}
        - 用户输入: "显示所有东西"
          你的输出: {{"command": "list", "arguments": []}}
        - 用户输入: "今天天气怎么样"
          你的输出: {{"command": "unknown", "arguments": ["无法处理天气查询"]}}
        """

        try:
            response = await ollama.AsyncClient().chat(
                model=self.model_name,
                messages=[
                    {'role': 'system', 'content': system_prompt},
                    {'role': 'user', 'content': user_input},
                ],
                format='json',  # 请求Ollama强制输出JSON格式
            )

            # response['message']['content'] 应该是一个JSON字符串
            content = response['message']['content']
            parsed_json = json.loads(content)

            # 验证解析后的JSON结构
            if 'command' in parsed_json and 'arguments' in parsed_json:
                return parsed_json
            else:
                return {'command': 'error', 'arguments': ['模型返回了无效的JSON结构。']}

        except json.JSONDecodeError:
            return {'command': 'error', 'arguments': ['无法解析模型的响应，它不是有效的JSON。']}
        except ollama.ResponseError as e:
            # 捕获Ollama的特定错误，例如模型未找到或服务未运行
            return {'command': 'error', 'arguments': [f"Ollama API 错误: {e.error}"]}
        except Exception as e:
            # 捕获其他通用异常，如网络问题
            return {'command': 'error', 'arguments': [f"处理时发生未知错误: {str(e)}"]}


class SimpleDBShell(cmd.Cmd):
    """
    一个简单的内存键值数据库交互式Shell，演示交互式终端界面的使用。
    """

    # 1. 自定义欢迎语和提示符
    lines = [
        "     欢迎使用 SimpleShell v1.0 (基于 cmd 库) ",
        " ",
        "这是一个简单的内存键值数据库，当程序退出时，所有数据将丢失",
        "在后续的代码中，这些功能将被保留",
        " ",
        "set <key> <value>  - 设置一个键值对 (值可以包含空格)",
        "get <key>          - 获取一个键对应的值",
        "delete <key>       - 删除一个键值对",
        "list               - 列出所有存储的键值对",
        "help               - 显示帮助信息",
        "exit / quit        - 退出程序"
    ]
    intro = create_aligned_intro(lines)
    prompt = "agent_db> "

    def __init__(self):
        """
        构造函数，初始化父类并创建一个用于存储数据的字典。
        """
        super().__init__()
        # 2. 在实例中维护状态（这里是一个字典作为内存数据库）
        self.data = {}
        # 添加自然语言处理器
        self.nlp = NaturalLanguageProcessor()

    # 3. 为每个命令实现一个 `do_<command>` 方法
    def do_set(self, arg):
        """
        设置一个键值对。用法: set <key> <value>
        值部分可以包含空格。
        """
        # `arg` 是用户在命令后输入的所有内容
        if not arg:
            print("❌ 错误: 'set' 命令需要一个键和一个值。")
            print("   用法: set <key> <value>")
            return

        # 使用 split(None, 1) 可以很好地处理值中包含空格的情况
        parts = arg.split(None, 1)
        if len(parts) < 2:
            print("❌ 错误: 请同时提供键和值。")
            print("   用法: set <key> <value>")
            return

        key, value = parts
        self.data[key] = value
        print(f"✅ 成功: '{key}' 已设置为 '{value}'")

    def do_get(self, arg):
        """
        根据键获取值。用法: get <key>
        """
        key = arg.strip()
        if not key:
            print("❌ 错误: 'get' 命令需要一个键。")
            print("   用法: get <key>")
            return

        value = self.data.get(key)
        if value is not None:
            print(f"➡️ {key} = {value}")
        else:
            print(f"🔍 未找到键: '{key}'")

    def do_delete(self, arg):
        """
        根据键删除一个键值对。用法: delete <key>
        """
        key = arg.strip()
        if not key:
            print("❌ 错误: 'delete' 命令需要一个键。")
            print("   用法: delete <key>")
            return

        if key in self.data:
            del self.data[key]
            print(f"✅ 成功: 键 '{key}' 已被删除。")
        else:
            print(f"🔍 未找到键: '{key}'")

    def do_list(self, arg):
        """
        列出所有存储的键值对。
        """
        if not self.data:
            print("ℹ️  数据库为空。")
            return

        print("📋 当前存储的所有数据:")
        max_key_len = max(len(k) for k in self.data.keys()) if self.data else 0
        for key, value in self.data.items():
            # 简单对齐，使输出更美观
            print(f"  - {key:<{max_key_len}} : {value}")

    # 4. 实现退出逻辑
    def do_exit(self, arg):
        """
        退出交互式Shell。
        """
        print("\n👋 感谢使用，再见！")
        return True  # 返回 True 会导致 cmdloop() 退出

    def do_quit(self, arg):
        """
        `quit` 是 `exit` 的别名。
        """
        return self.do_exit(arg)

    def do_EOF(self, arg):
        """
        处理 Ctrl+D 或Ctrl+Z 输入 (文件结束符)，实现优雅退出。
        """
        # `print()` 是为了在 ^D 后换行，使终端提示符在新的一行显示
        print()
        return self.do_exit(arg)

    # 5. (可选) 为特定命令提供更详细的帮助
    def help_set(self):
        print("\n".join([
            "命令: set <key> <value>",
            "功能: 创建或更新一个键值对。",
            "示例:",
            "  set name 'John Doe'  (设置一个带空格的值)",
            "  set age 30           (设置一个简单的值)"
        ]))

    # 6. (可选) 处理未知命令
    def default(self, line):
        """
        处理所有未匹配到 do_* 方法的输入，将其视为自然语言查询。
        """
        # 如果以'/'开头，则认为是未知命令，这与您原始代码的行为一致
        if line.startswith('/'):
            print(f"❌ 未知命令: '{line}'。输入 'help' 查看可用命令列表。")
            return

        # 定义一个异步函数来处理NLP逻辑
        async def process_natural_language():
            print(f"🤔 正在理解您的指令: '{line}'...")

            # 调用NLP处理器
            result = await self.nlp.process(line)

            command = result.get('command')
            arguments = result.get('arguments', [])

            print(f"🧠 模型解析结果: command='{command}', arguments={arguments}")

            # 根据解析结果执行相应的 do_* 方法
            if command == 'set':
                # 将参数列表合并成一个字符串，以匹配 do_set 的期望输入
                self.do_set(" ".join(arguments))
            elif command == 'get':
                self.do_get(" ".join(arguments))
            elif command == 'delete':
                self.do_delete(" ".join(arguments))
            elif command == 'list':
                self.do_list("")
            elif command == 'unknown':
                print(f"😕 抱歉，我无法理解您的请求。原因: {' '.join(arguments)}")
                print("   请尝试换一种方式描述，或者直接使用 'set', 'get' 等命令。")
            elif command == 'error':
                print(f"❌ 处理时发生错误: {' '.join(arguments)}")
                print("   请检查Ollama服务是否正在运行，以及模型是否已正确拉取。")
            else:
                print(f"⚠️ 模型返回了未知的命令: '{result}'")

        # 在同步的 default 方法中运行异步任务
        try:
            asyncio.run(process_natural_language())
        except Exception as e:
            print(f"❌ 在运行异步任务时发生严重错误: {e}")


def main():
    """
    程序主入口。
    """
    try:
        # 启动 cmd 循环
        SimpleDBShell().cmdloop()
    except KeyboardInterrupt:
        # 处理 Ctrl+C，实现优雅退出
        print("\n👋 检测到中断，程序退出。")
        sys.exit(0)


if __name__ == '__main__':
    main()